﻿// Copyright 2012 Jack Bulan (agentgeo)
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.IO;
using System.Windows.Forms;
using System.Diagnostics;
using System.Drawing;

using GemWorkshop.Formats.Images;

namespace GemWorkshop.Formats.Nitro
{
    public class ROMUnpacker
    {
        public string RomFileName = "", DirectoryName = "";
        private string LastDirectoryPath = "", ParentDirectoryPath = "";

        public string RomName, RomCode;

        public Dictionary<ushort, string> DirectoryPaths;
        public Dictionary<ushort, string> FilePaths;
        public Dictionary<ushort, ushort> FileNumbers;
        public Dictionary<ushort, ushort> ReverseFileNumbers;

        public AllocationTable AllocationTable;
        public NameTable NameTable;
        public ROMHeader Header;
        public ROMBanner Banner;
        public List<ARMOverlay> Overlays;

        private uint LastFileEnd = 0;
        public ushort FileCount = 0;

        private BinaryReader stream;

        #region Delegates & Events
        public delegate void ROMIsReadyD(TimeSpan timeTake);

        public delegate void DirReadyD(int DirID, int ParentID, string DirName, bool IsRoot);

        public delegate void FileReadyD(int FileID, int ParentID, string FileName);

        public event ROMIsReadyD ROMIsReady;

        public event DirReadyD DirReady;

        public event FileReadyD FileReady;
        #endregion

        public TimeSpan DoRom(string filePath, string outDirectory)
        {
            Stopwatch timer = Stopwatch.StartNew();

            if (Directory.Exists(outDirectory))
            {
                Directory.Delete(outDirectory, true);
            }
            Directory.CreateDirectory(outDirectory);
            DirectoryName = outDirectory;
            LastDirectoryPath = outDirectory;
            ParentDirectoryPath = LastDirectoryPath;

            Header = new ROMHeader();

            RomFileName = filePath;
            if (NameTable.IDs == null)
            {
                NameTable.IDs = new Dictionary<string, ushort>();
                NameTable.Names = new Dictionary<ushort, string>();
                AllocationTable.Offsets = new Dictionary<ushort, uint>();
                AllocationTable.Sizes = new Dictionary<ushort, uint>();
                DirectoryPaths = new Dictionary<ushort, string>();
                FilePaths = new Dictionary<ushort, string>();
                FileNumbers = new Dictionary<ushort, ushort>();
                ReverseFileNumbers = new Dictionary<ushort, ushort>();
                Overlays = new List<ARMOverlay>();
            }
            NameTable.IDs.Clear();
            NameTable.Names.Clear();
            AllocationTable.Offsets.Clear();
            AllocationTable.Sizes.Clear();
            DirectoryPaths.Clear();
            FileNumbers.Clear();
            FilePaths.Clear();
            ReverseFileNumbers.Clear();
            Overlays.Clear();

            stream = new BinaryReader(File.OpenRead(filePath));
            stream.BaseStream.Position = 0L;

            // Load the header
            #region Header
            stream.BaseStream.Position = 0L;

            Header.Title = Encoding.UTF8.GetString(stream.ReadBytes(12));
            Header.Code = Encoding.UTF8.GetString(stream.ReadBytes(4));
            Header.Maker = Encoding.UTF8.GetString(stream.ReadBytes(2));

            Header.DeviceCode = stream.ReadByte();
            Header.EncryptionSeed = stream.ReadByte();
            Header.FileLength = stream.ReadByte();
            stream.BaseStream.Seek(9L, SeekOrigin.Current);

            Header.Version = stream.ReadByte();
            Header.InternalFlags = stream.ReadByte();

            Header.Arm9Offset = stream.ReadUInt32();
            Header.Arm9Entry = stream.ReadUInt32();
            Header.Arm9Load = stream.ReadUInt32();
            Header.Arm9Length = stream.ReadUInt32();

            Header.Arm7Offset = stream.ReadUInt32();
            Header.Arm7Entry = stream.ReadUInt32();
            Header.Arm7Load = stream.ReadUInt32();
            Header.Arm7Length = stream.ReadUInt32();

            Header.FNTOffset = stream.ReadUInt32();
            Header.FNTSize = stream.ReadUInt32();
            Header.FATOffset = stream.ReadUInt32();
            Header.FATSize = stream.ReadUInt32();

            Header.Arm9OverlayOffset = stream.ReadUInt32();
            Header.Arm9OverlaySize = stream.ReadUInt32();
            Header.Arm7OverlayOffset = stream.ReadUInt32();
            Header.Arm7OverlaySize = stream.ReadUInt32();

            Header.NormalRegisterSettings = stream.ReadUInt32();
            Header.SecureRegisterSettings = stream.ReadUInt32();

            Header.BannerOffset = stream.ReadUInt32();

            Header.SecureAreaCRC = stream.ReadUInt16();
            Header.SecureTransferTimeout = stream.ReadUInt16();

            Header.Arm9AutoLoad = stream.ReadUInt32();
            Header.Arm7AutoLoad = stream.ReadUInt32();

            Header.SecureDisable = stream.ReadUInt64();

            Header.ROMLength = stream.ReadUInt32();
            Header.HeaderLength = stream.ReadUInt32();
            stream.BaseStream.Seek(56L, SeekOrigin.Current);

            Header.NintendoLogo = stream.ReadBytes(156);
            Header.NintendoLogoCRC = stream.ReadUInt16();
            Header.HeaderCRC = stream.ReadUInt16();
            #endregion

            // Load the folders and files.
            stream.BaseStream.Seek(Header.FNTOffset, SeekOrigin.Begin);
            NameTable.RawData = stream.ReadBytes((int)Header.FNTSize);
            //stream.BaseStream.Seek(0L, SeekOrigin.Begin);
            DoDirectory("root", 0xF000, 0);

            // Load the overlay files
            DoOverlays();

            stream.Close();
            stream.Dispose();

            ExtractBanner();

            // Extract Arm7, Arm9, and overlays
            File.WriteAllBytes(ParentDirectoryPath + "\\arm7.bin", ExtractSection(Header.Arm7Offset, Header.Arm7Length));
            File.WriteAllBytes(ParentDirectoryPath + "\\arm9.bin", ExtractSection(Header.Arm9Offset, Header.Arm9Length));
            File.WriteAllBytes(ParentDirectoryPath + "\\arm7overlay.bin", ExtractSection(Header.Arm7OverlayOffset, Header.Arm7OverlaySize));
            File.WriteAllBytes(ParentDirectoryPath + "\\arm9overlay.bin", ExtractSection(Header.Arm9OverlayOffset, Header.Arm9OverlaySize));

            timer.Stop();
            return timer.Elapsed;
        }

        private void DoDirectory(string DirName, ushort DirID, ushort Parent)
        {
            long iPos = stream.BaseStream.Position;
            stream.BaseStream.Position = (long)(Header.FNTOffset + (8 * (DirID & 0xFFF)));
            string tempDirPath = ParentDirectoryPath;

            uint EntryStart = stream.ReadUInt32();
            ushort EntryFileID = stream.ReadUInt16();
            ushort ParentID = stream.ReadUInt16();

            if (Parent == 0)
            {
                if (DirReady != null)
                    DirReady(0xF000, 0, "root", true);

                LastDirectoryPath = LastDirectoryPath + "\\root";
                ParentDirectoryPath = LastDirectoryPath;
                //MessageBox.Show("Creating directory: " + LastDirectoryPath);
                Directory.CreateDirectory(LastDirectoryPath);
                DirectoryPaths.Add(0xF000, LastDirectoryPath);
            }
            else
            {
                if (DirReady != null)
                    DirReady(DirID, Parent, DirName, false);

                LastDirectoryPath = ParentDirectoryPath + "\\" + DirName;
                ParentDirectoryPath = LastDirectoryPath;
                DirectoryPaths.Add(DirID, LastDirectoryPath);
                //MessageBox.Show("Creating directory: " + LastDirectoryPath);
                Directory.CreateDirectory(LastDirectoryPath);
            }

            stream.BaseStream.Position = (long)(Header.FNTOffset + EntryStart);

            ushort CurrentFile = EntryFileID;
            while (true)
            {
                byte EntryData = stream.ReadByte();
                byte NameLength = (byte)(EntryData & 0x7F);

                if (NameLength == 0) break;

                string EntryName = Encoding.UTF8.GetString(stream.ReadBytes(NameLength));

                if (EntryData > 0x7F)
                {
                    ushort SubTableID = stream.ReadUInt16();
                    DoDirectory(EntryName, SubTableID, DirID);
                }
                else
                {
                    DoFile(EntryName, CurrentFile, DirID);
                }
                CurrentFile++;
            }

            stream.BaseStream.Position = iPos;
            ParentDirectoryPath = tempDirPath;
        }

        private void DoFile(string FileName, ushort FileID, ushort Parent)
        {
            long iPos = stream.BaseStream.Position;

            stream.BaseStream.Position = (long)(Header.FATOffset + (FileID * 8));

            uint FATStart = stream.ReadUInt32();
            uint FATEnd = stream.ReadUInt32();

            if (FileReady != null)
                FileReady(FileID, Parent, FileName);

            NameTable.IDs[FileName] = FileID;
            NameTable.Names[FileID] = FileName;
            AllocationTable.Offsets[FileID] = FATStart;
            AllocationTable.Sizes[FileID] = FATEnd - FATStart;
            FilePaths[FileCount] = DirectoryPaths[Parent] + "\\" + FileName;
            FileNumbers[FileID] = FileCount;
            ReverseFileNumbers[FileCount] = FileID;
            //FilePaths.Add(FileID, DirectoryPaths[Parent] + "\\" + FileName);

            stream.BaseStream.Position = (long)AllocationTable.Offsets[FileID];
            byte[] buffer = stream.ReadBytes((int)AllocationTable.Sizes[FileID]);
            File.WriteAllBytes(DirectoryPaths[Parent] + "\\" + FileName, buffer);

            if (FATEnd > LastFileEnd)
                LastFileEnd = FATEnd;

            FileCount++;

            stream.BaseStream.Position = iPos;
        }

        private void DoOverlays()
        {
            string overlayDir = DirectoryName + "\\overlay";
            Directory.CreateDirectory(overlayDir);
            
            // Arm9 Overlay
            if (Header.Arm9OverlaySize > 0)
            {
                stream.BaseStream.Seek(Header.Arm9OverlayOffset, SeekOrigin.Begin);
                int files = (int)Header.Arm9OverlaySize / 32;

                ARMOverlay overlay = new ARMOverlay();
                for (int i = 0; i < files; i++)
                {
                    // Read from the table
                    overlay = new ARMOverlay();
                    overlay.ID = stream.ReadUInt32();
                    overlay.RAMOffset = stream.ReadUInt32();
                    overlay.RAMSize = stream.ReadUInt32();
                    overlay.BSSSize = stream.ReadUInt32();
                    overlay.InitStart = stream.ReadUInt32();
                    overlay.InitEnd = stream.ReadUInt32();
                    overlay.FileID = stream.ReadUInt32();
                    stream.BaseStream.Seek(4L, SeekOrigin.Current); // Reserved

                    // Read the file
                    long iPos = stream.BaseStream.Position;
                    uint fatOffset = overlay.FileID * 8;
                    stream.BaseStream.Seek(Header.FATOffset + fatOffset, SeekOrigin.Begin);

                    uint start = stream.ReadUInt32();
                    uint end = stream.ReadUInt32();
                    overlay.Offest = start;
                    overlay.Size = end - start;
                    //stream.BaseStream.Position = iPos;

                    // Write to a file
                    stream.BaseStream.Seek(overlay.Offest, SeekOrigin.Begin);
                    File.WriteAllBytes(overlayDir + "\\overlay" + i.ToString() + ".bin", stream.ReadBytes((int)overlay.Size));

                    stream.BaseStream.Seek(iPos, SeekOrigin.Begin); // Back to our start

                    Overlays.Add(overlay);
                }
            }
        }

        public byte[] ExtractFile(ushort FileID)
        {
            stream = new BinaryReader(File.OpenRead(RomFileName));

            stream.BaseStream.Position = (long)AllocationTable.Offsets[FileID];
            byte[] buffer = stream.ReadBytes((int)AllocationTable.Sizes[FileID]);

            stream.Close();
            stream.Dispose();

            return buffer;
        }

        public byte[] ExtractSection(uint offset, uint length)
        {
            stream = new BinaryReader(File.OpenRead(RomFileName));

            stream.BaseStream.Position = (long)offset;
            byte[] buffer = stream.ReadBytes((int)length);

            stream.Close();
            stream.Dispose();

            return buffer;
        }

        private void ExtractBanner()
        {
            BinaryReader br = new BinaryReader(File.OpenRead(RomFileName));

            br.BaseStream.Seek((long)Header.BannerOffset, SeekOrigin.Begin);
            Banner = new ROMBanner();

            // Load its values
            Banner.Version = br.ReadUInt16();
            Banner.CRC16 = br.ReadUInt16();
            br.BaseStream.Seek(28L, SeekOrigin.Current);
            Banner.TileData = br.ReadBytes(512);
            Banner.Palette = br.ReadBytes(32);
            Banner.JapaneseTitle = Encoding.Unicode.GetString(br.ReadBytes(256));
            Banner.EnglishTitle = Encoding.Unicode.GetString(br.ReadBytes(256));
            Banner.FrenchTitle = Encoding.Unicode.GetString(br.ReadBytes(256));
            Banner.GermanTitle = Encoding.Unicode.GetString(br.ReadBytes(256));
            Banner.ItalianTitle = Encoding.Unicode.GetString(br.ReadBytes(256));
            Banner.SpanishTitle = Encoding.Unicode.GetString(br.ReadBytes(256));

            Banner.Image = new Bitmap(32, 32);
            
            // Now, draw it UP!
            byte[] tiles = BitsConverter.BytesToBit4(Banner.TileData);
            Color[] palette = Actions.BGR555ToColor(Banner.Palette);
            int i = 0;
            for (int hi = 0; hi < 4; hi++)
            {
                for (int wi = 0; wi < 4; wi++)
                {
                    for (int h = 0; h < 8; h++)
                    {
                        for (int w = 0; w < 8; w++)
                        {
                            Banner.Image.SetPixel(w + wi * 8, h + hi * 8, palette[tiles[i]]);
                            i++;
                        }
                    }
                }
            }


            br.Close();
            br.Dispose();
        }
    }

    public struct ROMHeader
    {
        // It follows this format exactly
        public string Title, Code, Maker;
        public byte DeviceCode, EncryptionSeed;
        public byte FileLength;
        // 9 bytes reserved (00)
        public byte Version, InternalFlags;
        public uint Arm9Offset, Arm9Entry, Arm9Load, Arm9Length;
        public uint Arm7Offset, Arm7Entry, Arm7Load, Arm7Length;
        public uint FNTOffset, FNTSize;
        public uint FATOffset, FATSize;
        public uint Arm9OverlayOffset, Arm9OverlaySize;
        public uint Arm7OverlayOffset, Arm7OverlaySize;
        public uint NormalRegisterSettings, SecureRegisterSettings;
        public uint BannerOffset;
        public ushort SecureAreaCRC, SecureTransferTimeout;
        public uint Arm9AutoLoad, Arm7AutoLoad;
        public ulong SecureDisable;
        public uint ROMLength;
        public uint HeaderLength;
        // 56 bytes reserved (00)
        public byte[] NintendoLogo; // 156 bytes
        public ushort NintendoLogoCRC, HeaderCRC;
        //public byte[] DebuggerReserved, ConfigurationSettings;
        // 160 bytes reserved (00)
    }

    public struct ROMBanner
    {
        public ushort Version, CRC16;
        // 28 bytes reserved
        public byte[] TileData, Palette;
        public string JapaneseTitle, EnglishTitle, FrenchTitle;
        public string GermanTitle, ItalianTitle, SpanishTitle;

        public Bitmap Image;
    }

    public struct NameTable
    {
        public byte[] RawData; // Used for saving purposes.
        public Dictionary<ushort, string> Names;
        public Dictionary<string, ushort> IDs;
    }

    public struct AllocationTable
    {
        public Dictionary<ushort, uint> Sizes, Offsets;
    }

    public struct ARMOverlay
    {
        public uint ID;
        public uint RAMOffset, RAMSize;
        public uint BSSSize;
        public uint InitStart, InitEnd;
        public uint FileID;
        // 4 bytes reserved

        public uint Offest, Size;

        public bool ARM9;
        public string Path;
    }
}
